#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
"""This script will execute `pip-compile` to generate a variety of environment specifications.

An environment specification is a requirements.txt file with fully pinned versions for all
dependencies and transitive dependencies. Each CI job uses a specific environment specification
generated by this script.


## Setting up factorized dependencies
The deps/ subdirectory contains "factorized" lists of dependencies. For example, `runtime.txt`
contains the loosest possible specification of Python packages that must be installed for
the code to run. Additional files specify additional dependencies: for example, `pytest.txt`
contains the additional dependencies to run the `pytest` CI job.

The PlatformRecipes in this file provide a set of factorized dependency (file-)names
for each full environment. For example, the `pytest` environment is composed of the `runtime`
and `pytest` factorized dependency files.


## Consistency and `dev.env.txt`
Before generating any of the requested environments, we create a special environment named
`dev.env.txt` for each PlatformRecipe. This contains all packages named in any of the factorized
dependencies. Remember that `pip-compile` will find a set of version pins such that all
dependencies' requirements are satisfied. This environment provides a fully consistent set
of version pins. We use it as a pip "constraint" for the other environments. That means when
different environments share dependencies (within a platform), they will always be the same version.
"""

import os
import subprocess
from argparse import ArgumentParser
from dataclasses import dataclass, field
from typing import *


@dataclass
class EnvRecipe:
    """Make one environment with `deps`"""

    deps: Set[str]
    addtl_constraints: Set[str] = field(default_factory=set)


@dataclass
class PlatformRecipe:
    """Make all named `envs`."""

    envs: Dict[str, EnvRecipe]
    env_out_dir: str = "envs"


PLATFORMS = {
    "default": PlatformRecipe(
        {
            "format": EnvRecipe({"runtime", "format"}),
            "pylint": EnvRecipe({"runtime", "pylint"}),
            "pytest": EnvRecipe({"runtime", "pytest"}),
            "pytest-extra": EnvRecipe({"runtime", "resource_estimates_runtime", "pytest"}),
            "mypy": EnvRecipe({"runtime", "mypy"}),
            "pip-tools": EnvRecipe({"pip-tools"}),
        }
    ),
    # The following is intended to be run on an older version of python and
    # includes the additional deps/oldest-versions.txt constraint file.
    "max_compat": PlatformRecipe(
        {
            "pytest-max-compat": EnvRecipe(
                {"runtime", "pytest"}, addtl_constraints={"oldest-versions"}
            )
        },
        env_out_dir="max_compat",
    ),
}


def run(*args):
    """Run a command using `subprocess`."""
    return subprocess.run(*args, check=True)


def pip_compile(env_name: str, env_recipe: EnvRecipe, env_out_dir: str, constrain=True):
    """Run `pip-compile` to create the named environment."""
    dep_args = [f"deps/{dep_name}.txt" for dep_name in env_recipe.deps]
    dep_args += [f"--constraint=deps/{cons_name}.txt" for cons_name in env_recipe.addtl_constraints]
    if constrain:
        dep_args.append(f"--constraint={env_out_dir}/dev.env.txt")

    run(
        [
            "pip-compile",
            f"--output-file={env_out_dir}/{env_name}.env.txt",
            "--resolver=backtracking",
        ]
        + dep_args
    )


def get_dev_env_recipe(pr: PlatformRecipe) -> EnvRecipe:
    """Add up all the factorized dependencies."""
    all_deps = set()
    all_addtl_constraints = set()
    for env_recipe in pr.envs.values():
        all_deps |= env_recipe.deps
        all_addtl_constraints |= env_recipe.addtl_constraints
    return EnvRecipe(all_deps, all_addtl_constraints)


def make_platform_envs(pr: PlatformRecipe):
    os.makedirs(pr.env_out_dir, exist_ok=True)

    # Pip compile the full dev environment
    pip_compile("dev", get_dev_env_recipe(pr), pr.env_out_dir, constrain=False)

    # Pip compile the environments
    for env_name, env_recipe in pr.envs.items():
        pip_compile(env_name, env_recipe, pr.env_out_dir)


def parse():
    """Parse command line arguments."""
    parser = ArgumentParser()
    parser.add_argument("--platform", default="default")
    args = parser.parse_args()
    try:
        platform = PLATFORMS[args.platform]
    except KeyError:
        raise ValueError(f"Unknown platform {args.platform}")

    return make_platform_envs(platform)


if __name__ == "__main__":
    parse()
